Stochastic SIR in continuous time using DifferentialEquations.jl
Simon Frost
using DifferentialEquations
using SimpleDiffEq
using Random
using Plots
using BenchmarkTools
function infection_rate(u,p,t)
(S,I,R) = u
(β,γ) = p
N = S+I+R
β*S*I/N
end
function infection!(integrator)
integrator.u[1] -= 1
integrator.u[2] += 1
end
infection_jump = ConstantRateJump(infection_rate,infection!)
DiffEqJump.ConstantRateJump{typeof(Main.WeaveSandBox2.infection_rate),typeo
f(Main.WeaveSandBox2.infection!)}(Main.WeaveSandBox2.infection_rate, Main.W
eaveSandBox2.infection!)
function recovery_rate(u,p,t)
(S,I,R) = u
(β,γ) = p
γ*I
end
function recovery!(integrator)
integrator.u[2] -= 1
integrator.u[3] += 1
end
recovery_jump = ConstantRateJump(recovery_rate,recovery!)
DiffEqJump.ConstantRateJump{typeof(Main.WeaveSandBox2.recovery_rate),typeof
(Main.WeaveSandBox2.recovery!)}(Main.WeaveSandBox2.recovery_rate, Main.Weav
eSandBox2.recovery!)
tspan = (0.0,50.0)
u0 = [999.0,1.0,0.0]
p = [0.5,0.25]
Random.seed!(1234)
Random.MersenneTwister(UInt32[0x000004d2], Random.DSFMT.DSFMT_state(Int32[-
1393240018, 1073611148, 45497681, 1072875908, 436273599, 1073674613, -20437
16458, 1073445557, -254908435, 1072827086 … -599655111, 1073144102, 36765
5457, 1072985259, -1278750689, 1018350124, -597141475, 249849711, 382, 0]),
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 … 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0, 0.0, 0.0], UInt128[0x00000000000000000000000000000000,
0x00000000000000000000000000000000, 0x00000000000000000000000000000000, 0x0
0000000000000000000000000000000, 0x00000000000000000000000000000000, 0x0000
0000000000000000000000000000, 0x00000000000000000000000000000000, 0x0000000
0000000000000000000000000, 0x00000000000000000000000000000000, 0x0000000000
0000000000000000000000 … 0x00000000000000000000000000000000, 0x0000000000
0000000000000000000000, 0x00000000000000000000000000000000, 0x0000000000000
0000000000000000000, 0x00000000000000000000000000000000, 0x0000000000000000
0000000000000000, 0x00000000000000000000000000000000, 0x0000000000000000000
0000000000000, 0x00000000000000000000000000000000, 0x0000000000000000000000
0000000000], 1002, 0)
prob = DiscreteProblem(u0,tspan,p)
prob_sir_jump = JumpProblem(prob,Direct(),infection_jump,recovery_jump)
sol_sir_jump = solve(prob_sir_jump,FunctionMap())
retcode: Success
Interpolation: left-endpoint piecewise constant
t: 3151-element Array{Float64,1}:
0.0
0.0
1.0
2.0
2.215004782899086
2.215004782899086
2.6506894802326073
2.6506894802326073
3.6506894802326073
3.9639892224020175
⋮
48.19448263110623
49.19448263110623
49.24586855215085
49.24586855215085
49.65703941639901
49.65703941639901
49.671067280710226
49.671067280710226
50.0
u: 3151-element Array{Array{Float64,1},1}:
[999.0, 1.0, 0.0]
[999.0, 1.0, 0.0]
[999.0, 1.0, 0.0]
[999.0, 1.0, 0.0]
[999.0, 1.0, 0.0]
[998.0, 2.0, 0.0]
[998.0, 2.0, 0.0]
[998.0, 1.0, 1.0]
[998.0, 1.0, 1.0]
[998.0, 1.0, 1.0]
⋮
[212.0, 7.0, 781.0]
[212.0, 7.0, 781.0]
[212.0, 7.0, 781.0]
[212.0, 6.0, 782.0]
[212.0, 6.0, 782.0]
[212.0, 5.0, 783.0]
[212.0, 5.0, 783.0]
[211.0, 6.0, 783.0]
[211.0, 6.0, 783.0]
plot(sol_sir_jump,vars=[(0,1),(0,2),(0,3)])
@benchmark solve(prob_sir_jump,FunctionMap())
BenchmarkTools.Trial:
memory estimate: 13.48 KiB
allocs estimate: 116
--------------
minimum time: 9.599 μs (0.00% GC)
median time: 128.399 μs (0.00% GC)
mean time: 409.863 μs (6.08% GC)
maximum time: 10.489 ms (91.22% GC)
--------------
samples: 10000
evals/sample: 1